#pragma once

#include <cstdint>
#include <future>
#include <vector>
#include <queue>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <functional>

#ifndef NETWORK_TRAFFIC_PROCESSING_V2_ThreadPoolV2
#define NETWORK_TRAFFIC_PROCESSING_V2_ThreadPoolV2
class ThreadPool {
private:
    // storage for threads and tasks
    std::vector<std::thread> threads;
    std::queue<std::function<void(void)>> tasks;

    // primitives for signaling
    std::mutex mutex;
    std::condition_variable cv;

    // the state of the thread, pool
    bool stop_pool;
    uint32_t active_threads;
    const uint64_t capacity;

    std::atomic<size_t> task_cnt{ 0 };
    std::atomic<size_t> task_cnt_wait{ 0 };
    std::atomic<size_t> submit{ 0 };

    // custom task factory
    template <typename Func, typename ... Args, typename Rtrn = typename std::result_of<Func(Args...)>::type>
    auto make_task(Func && func, Args && ...args) -> std::packaged_task<Rtrn(void)>
    {
        auto aux = std::bind(std::forward<Func>(func), std::forward<Args>(args)...);
        return std::packaged_task<Rtrn(void)>(aux);
    }

    // will be executed before execution of a task
    void before_task_hook() { active_threads++; }

    // will be executed after execution of a task
    void after_task_hook() { active_threads--; }

public:
    ThreadPool(uint64_t capacity_) :
            stop_pool(false),     // pool is running
            active_threads(0),    // no work to be done
            capacity(capacity_)   // remember size
    {
        // this function is executed by the threads
        auto wait_loop = [this]() -> void
        {
            // wait forever
            while (true) {
                // this is a placeholder task
                std::function<void(void)> task;
                {
                    // lock this section for waiting
                    std::unique_lock<std::mutex> unique_lock(mutex);

                    // actions must be performed on wake-up if (i) the thread pool
                    // has been stopped or (ii) there are still tasks to be processed
                    auto predicate = [this]() -> bool {
                        return  (stop_pool) || !(tasks.empty());
                    };

                    task_cnt_wait++;
                    // wait to be waken up on aforementioned conditions
                    cv.wait(unique_lock, predicate);

                    // exit if thread pool stopped and no tasks to be performed
                    if (stop_pool && tasks.empty())
                        return;

                    submit++;
                    // else extract task from queue
                    task = std::move(tasks.front());
                    tasks.pop();
                    before_task_hook();
                } // here we release the lock

                // execute the task in parallel
                task();
                task_cnt++;
                {   // adjust the thread counter
                    std::lock_guard<std::mutex> lock_guard(mutex);
                    after_task_hook();
                } // here we release the lock
            }
        };

        // initially spawn capacity many threads
        for (uint64_t id = 0; id < capacity; id++)
            threads.emplace_back(wait_loop);
    }

    ~ThreadPool()
    {

        { // acquire a scoped lock
            std::lock_guard<std::mutex> lock_guard(mutex);

            // and subsequently alter
            // the global state to stop
            stop_pool = true;
        } // here we release the lock

        // signal all threads
        cv.notify_all();

        // finally join all threads
        for (auto& thread : threads)
            thread.join();
    }

    template <typename Func, typename ... Args, typename Pair = Func(Args...), typename Rtrn = typename std::result_of<Pair>::type>
    void enqueue(Func && func, Args && ... args) //-> std::future<Rtrn>
    {
        // create the task, get the future and wrap task in a shared pointer
        auto task = make_task(func, args...);
        //auto future = task.get_future();
        auto task_ptr = std::make_shared<decltype(task)>(std::move(task));

        {   // lock the scope
            std::lock_guard<std::mutex> lock_guard(mutex);

            // you cannot reuse pool after being stopped
            if (stop_pool) throw std::runtime_error("enqueue on stopped ThreadPool");

            // wrap the task in a generic void function void -> void
            auto payload = [task_ptr]() -> void {
                // basically call task()
                task_ptr->operator()();
            };

            // append the task to the queue
            tasks.emplace(payload);
        }

        // tell one thread to wake-up
        cv.notify_one();

        //return future;
    }

    size_t get_active_threads() { return active_threads; }
    size_t get_queue_size() { return tasks.size(); }
    size_t get_cnt_size() { return submit.load(); }
    size_t get_task_cnt() { return task_cnt.load(); }
    size_t get_task_cnt_wait() { return task_cnt_wait.load(); }

};

#endif // NETWORK_TRAFFIC_PROCESSING_V2_ThreadPoolV2

